﻿#include "opencv_raw.h"

void raw_thread::read_raw_5files(const std::vector<cv::Mat>& raw_frames) {
	// 1. 多帧ORB对齐
	std::vector<cv::Mat> aligned_frames;
	aligned_frames.push_back(raw_frames[0]);
	for (size_t i = 1; i < raw_frames.size(); ++i) {
		cv::Mat ref_8u, frame_8u;
		cv::normalize(raw_frames[0], ref_8u, 0, 255, cv::NORM_MINMAX, CV_8UC1);
		cv::normalize(raw_frames[i], frame_8u, 0, 255, cv::NORM_MINMAX, CV_8UC1);

		cv::Mat aligned_8u, H;
		aligned_8u = alignWithORB(ref_8u, frame_8u, H);
		if (H.empty()) {
			qWarning() << "第" << i << "帧对齐失败，使用原图";
			aligned_frames.push_back(raw_frames[i].clone());
			continue;
		}

		cv::Mat aligned_16u;
		cv::warpPerspective(raw_frames[i], aligned_16u, H, raw_frames[0].size(),
			cv::INTER_LINEAR, cv::BORDER_CONSTANT, cv::Scalar(0));
		aligned_frames.push_back(aligned_16u);
		qDebug() << "第" << i << "帧16位对齐完成";
	}

	// 2. 多帧平均去噪
	cv::Mat avg_result = multiFrameAverage(aligned_frames);
	cv::imwrite("1_avg_result.png", avg_result);

	// 3. 16位转8位进行图像处理
	cv::Mat display_8u;
	double min_val, max_val;
	cv::minMaxLoc(avg_result, &min_val, &max_val);
	if (max_val - min_val < 30000) {
		cv::normalize(avg_result, display_8u, 0, 255, cv::NORM_MINMAX, CV_8UC1);
	}
	else {
		avg_result.convertTo(display_8u, CV_8UC1, 255.0 / 65535.0);
	}
	cv::imwrite("2_display_8u.png", display_8u);

	// 4. CLAHE对比度增强
	cv::Mat clahe_img;
	cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(2.0, cv::Size(4, 4));
	clahe->apply(display_8u, clahe_img);
	cv::imwrite("3_clahe.png", clahe_img);

	// 5. Sobel边缘提取
	cv::Mat sobel_x, sobel_y, edge;
	cv::Sobel(clahe_img, sobel_x, CV_16S, 1, 0, 3);
	cv::Sobel(clahe_img, sobel_y, CV_16S, 0, 1, 3);
	cv::convertScaleAbs(sobel_x, sobel_x);
	cv::convertScaleAbs(sobel_y, sobel_y);
	cv::addWeighted(sobel_x, 0.5, sobel_y, 0.5, 0, edge);
	cv::imwrite("4_Sobel.png", edge);

	// 6. 形态学梯度增强
	cv::Mat morph_grad;
	cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
	cv::morphologyEx(clahe_img, morph_grad, cv::MORPH_GRADIENT, kernel);
	cv::imwrite("5_morph_grad.png", morph_grad);

	// 7. 多尺度融合
	cv::Mat fused;
	cv::addWeighted(clahe_img, 0.6, edge, 0.15, 0, fused);
	cv::addWeighted(fused, 1.0, morph_grad, 0.15, 0, fused);
	cv::imwrite("6_fused.png", fused);

	// 8. 伽马亮度增强
	double gamma = 0.6;
	cv::Mat gamma_corrected = gammaCorrection(fused, gamma);
	cv::imwrite("7_gamma_corrected.png", gamma_corrected);

	// 9. 转换为3通道（适配ESPCN输入）
	cv::Mat src3c;
	cv::cvtColor(gamma_corrected, src3c, cv::COLOR_GRAY2BGR);

	// 10. ESPCN超分
	cv::Mat sr_img;
	bool sr_success = true;
	std::string model_path = "ESPCN_x4.pb";  // 模型路径
	try {
		cv::dnn_superres::DnnSuperResImpl sr;
		sr.readModel(model_path);
		sr.setModel("espcn", 4);
		sr.upsample(src3c, sr_img);
		cv::cvtColor(sr_img, sr_img, cv::COLOR_BGR2GRAY); // 转回单通道
		qDebug() << "ESPCN超分后尺寸:" << sr_img.cols << "x" << sr_img.rows;

		if (!cv::imwrite("8_ESPCN_output.png", sr_img)) {
			qWarning() << "超分结果保存失败";
		}
	}
	catch (const cv::Exception& e) {
		qCritical() << "超分失败：" << e.what();
		sr_img = gamma_corrected.clone();
		sr_success = false;
	}

	if (!sr_success) {
		emit send_image(gamma_corrected);
		return;
	}

	// 发送最终结果在label中显示
	emit send_image(sr_img);
}

// ORB特征对齐（输入8位图像，返回对齐后的8位图像和单应性矩阵H）
cv::Mat raw_thread::alignWithORB(const cv::Mat& ref_8u, const cv::Mat& frame_8u, cv::Mat& H) {
	H = cv::Mat(); // 初始化H为空
	if (ref_8u.empty() || frame_8u.empty()) {
		qCritical() << "ORB对齐失败：输入图像为空";
		return frame_8u.clone();
	}

	// 初始化ORB检测器
	cv::Ptr<cv::ORB> orb = cv::ORB::create(8000, 1.2f, 8);

	// 提取特征点和描述符
	std::vector<cv::KeyPoint> ref_kp, frame_kp;
	cv::Mat ref_desc, frame_desc;
	orb->detectAndCompute(ref_8u, cv::noArray(), ref_kp, ref_desc);
	orb->detectAndCompute(frame_8u, cv::noArray(), frame_kp, frame_desc);

	qDebug() << "ORB特征点数量（参考帧：" << ref_kp.size() << "，当前帧：" << frame_kp.size() << "）";

	// 检查特征点数量
	if (ref_kp.size() < 10 || frame_kp.size() < 10) {
		qWarning() << "特征点不足，返回原图";
		return frame_8u.clone();
	}

	// 匹配描述符
	cv::BFMatcher matcher(cv::NORM_HAMMING);
	std::vector<cv::DMatch> matches;
	matcher.match(ref_desc, frame_desc, matches);

	// 筛选优质匹配
	if (matches.size() < 10) {
		qWarning() << "有效匹配点不足（" << matches.size() << "），返回原图";
		return frame_8u.clone();
	}
	std::sort(matches.begin(), matches.end(), [](const cv::DMatch& a, const cv::DMatch& b) {
		return a.distance < b.distance;
		});
	int keep = std::max(10, (int)(matches.size() * 0.2));
	matches.resize(keep);

	// 提取匹配点坐标
	std::vector<cv::Point2f> ref_pts, frame_pts;
	for (const auto& m : matches) {
		ref_pts.push_back(ref_kp[m.queryIdx].pt);
		frame_pts.push_back(frame_kp[m.trainIdx].pt);
	}

	// 计算单应性矩阵H
	H = cv::findHomography(frame_pts, ref_pts, cv::RANSAC, 5.0);
	if (H.empty()) {
		qWarning() << "单应性矩阵计算失败，返回原图";
		return frame_8u.clone();
	}

	// 对齐8位图像
	cv::Mat aligned_8u;
	cv::warpPerspective(
		frame_8u,
		aligned_8u,
		H,
		ref_8u.size(),
		cv::INTER_LINEAR,
		cv::BORDER_CONSTANT,
		cv::Scalar(0)
	);

	return aligned_8u;
}

// 多帧平均
cv::Mat raw_thread::multiFrameAverage(const std::vector<cv::Mat>& aligned_frames) {
	if (aligned_frames.empty()) {
		qCritical() << "多帧平均失败：输入帧为空";
		return cv::Mat();
	}

	// 校验帧尺寸和类型
	int rows = aligned_frames[0].rows;
	int cols = aligned_frames[0].cols;
	for (const auto& frame : aligned_frames) {
		if (frame.rows != rows || frame.cols != cols || frame.type() != CV_16UC1) {
			qCritical() << "多帧平均失败：帧尺寸或类型不匹配";
			return cv::Mat();
		}
	}

	// 32位整数累加避免溢出
	cv::Mat sum_mat(rows, cols, CV_32SC1, cv::Scalar(0));
	for (const auto& frame : aligned_frames) {
		cv::Mat frame_32s;
		frame.convertTo(frame_32s, CV_32SC1);
		sum_mat += frame_32s;
	}

	// 计算平均值并转回16位
	cv::Mat avg_mat;
	sum_mat.convertTo(avg_mat, CV_16UC1, 1.0 / aligned_frames.size(), 0.5); // 四舍五入

	return avg_mat;
}

// 伽马亮度增强
cv::Mat raw_thread::gammaCorrection(const cv::Mat& src, double gamma) {
	cv::Mat lut(1, 256, CV_8UC1);
	uchar* ptr = lut.ptr();
	for (int i = 0; i < 256; ++i) {
		ptr[i] = cv::saturate_cast<uchar>(255.0 * pow(i / 255.0, gamma));
	}
	cv::Mat dst;
	cv::LUT(src, lut, dst);
	return dst;
}